package com.wlz.config;

import com.wlz.utils.TraceIdUtil;
import org.springframework.http.HttpRequest;
import org.springframework.http.client.ClientHttpRequestExecution;
import org.springframework.http.client.ClientHttpRequestInterceptor;
import org.springframework.http.client.ClientHttpResponse;

import java.io.IOException;

/**
 * @author wlz
 */
public class RestTemplateTraceIdInterceptor implements ClientHttpRequestInterceptor {


    @Override
    public ClientHttpResponse intercept(HttpRequest request, byte[] body,
                                        ClientHttpRequestExecution clientHttpRequestExecution) throws IOException {
        String traceId= TraceIdUtil.getTraceId();
        if(traceId!=null) {
            request.getHeaders().set(TraceIdUtil.TRACE_ID, traceId);
        } else {
            request.getHeaders().set(TraceIdUtil.TRACE_ID, TraceIdUtil.generateTraceId());
        }
        return clientHttpRequestExecution.execute(request, body);
    }

}